'''
https://leetcode.cn/problems/shortest-path-to-get-all-keys/description/
'''
from collections import deque
from typing import List
import numpy as np

class Solution:
    # 分层图 bfs
    # 分层图: 一个点位置+状态 为图上的一个顶点
    #        for example:(0,0,0), (0,0,1), (0,0,2), (0,0,3) 为都在(0,0)点，但是不同状态，也算不同的点
    def shortestPathAllKeys(self, grid: List[str]) -> int:
        start = (0, 0, 0)
        m, n = len(grid), len(grid[0])
        target_state = 0
        for x in range(m):
            for y in range(n):
                if grid[x][y] == '@':
                    start = (x, y, 0)
                if grid[x][y].islower():
                    target_state |= (1 << (ord(grid[x][y]) - ord('a')))

        q = deque([start])
        visited = np.full((m, n , target_state), False, dtype=bool)
        visited[start] = True
        ans = 0
        while q:
            for _ in range(len(q)):
                x, y, state = q.popleft()
                for nx, ny in [(x+1, y), (x-1, y), (x, y+1), (x, y-1)]:
                    if 0 <= nx < m and 0 <= ny < n and grid[nx][ny] != '#':
                        n_state = state
                        if grid[nx][ny].islower():
                            n_state = state | (1 << (ord(grid[nx][ny]) - ord('a')))
                        elif grid[nx][ny].isupper():
                            if not (state & (1 << (ord(grid[nx][ny]) - ord('A')))):
                                continue
                        if n_state == target_state:
                            return ans + 1
                        if not visited[nx, ny, n_state]:
                            visited[nx, ny, n_state] = True
                            q.append((nx, ny, n_state))
            ans += 1
        return -1

grid = ["@.a..","###.#","b.A.B"]
print(Solution().shortestPathAllKeys(grid))
